import numpy as np
from sklearn.gaussian_process import GaussianProcessRegressor
from sklearn.gaussian_process.kernels import ConstantKernel, RBF
import matplotlib.pyplot as plt
from error_file import NotFittedError


class GaussProcessRegressor(object):
    def __init__(self):
        kernel_ = ConstantKernel(constant_value=0.2, constant_value_bounds=(1e-4, 1e4)) * \
                 RBF(length_scale=0.5, length_scale_bounds=(1e-4, 1e4))
        self.gpr = GaussianProcessRegressor(kernel=kernel_, n_restarts_optimizer=2)
        self._is_fitted = False

    def init_data(self, x_train, y_train):
        """
        param: x_train
        param: y_train
        return:
        """
        (m, n) = np.shape(x_train)
        tem_m = np.shape(y_train)
        try:
            if m != tem_m[0]:
                raise ValueError
            self.x_train = x_train
            self.y_train = y_train
        except ValueError:
            print("数据出错，x,y样本不一致")
            return

    def init_set(self):
        pass

    def fit(self, x_train, y_train):
        self.gpr.fit(x_train, y_train)
        self._is_fitted = True

    def predict(self, x_test, return_std=False, return_cov=False):
        if self._is_fitted:
            mu_, cov_ = self.gpr.predict(x_test, return_std, return_cov)
            return mu_, cov_
        else:
            raise NotFittedError("fit before predict!")

    def get_mse(self, y_test):
        pass

    def get_rmse(self, y_test):
        pass

    def get_mae(self, y_test):
        pass

    def get_score(self, x_test, y_test):
        pass


def target(x):
    return np.exp(-(x - 2) ** 2) + np.exp(-(x - 6) ** 2 / 10) + 1 / (x ** 2 + 1)


if __name__ == '__main__':
    # 训练和测试数据
    X_train = np.array([-1, 2, 4, 5, 9]).reshape(-1, 1)
    y_train = target(X_train)
    X_test = np.linspace(-2, 10, 10000).reshape(-1, 1)

    gpr = GaussProcessRegressor()
    gpr.fit(X_train, y_train)
    mu, cov = gpr.predict(X_test, return_cov=True)
    y_test = mu.ravel()
    uncertainty = 1.96 * np.sqrt(np.diag(cov))
    plt.figure()
    plt.title("l=%.1f sigma_f=%.1f" % (gpr.gpr.kernel_.k2.length_scale, gpr.gpr.kernel_.k1.constant_value))
    plt.fill_between(X_test.ravel(), y_test + uncertainty, y_test - uncertainty, alpha=0.1)
    plt.plot(X_test, y_test, label="predict")
    plt.scatter(X_train, y_train, label="train", c="red", marker="D")
    plt.legend()
    plt.show()




